#ifndef _INTEGERMAP_H_
#define _INTEGERMAP_H_

#include "util.h"
#include "temparena.h"

// Split nodes into this number of nodes when they become nonterminal
#define ARITY 64
// The log of ARITY
#define LG_ARITY 6

// Determines for a Key in a given node, which child it belongs under
// Determined by a digit in the base ARITY representation of Key
// lg_range represents the range of the node you are currently in
// which is why this grabs the bits [lg_range - LG_ARITY, lg_range)
inline uint64 TopBits(uint64 k, uint64 lg_range) {
  return (k >> (lg_range - LG_ARITY)) & (ARITY - 1);
}

template <class V> 
struct IntegerMap {
  typedef pair<uint64, V> KV;  

  struct Node;

  // all iterators are invalidated on map modification.
  struct iterator {
    // end iterator points to root with a pos of -1
    Node * n_;
    int64 pos_;
  iterator(Node *n, int64 pos) :n_(n), pos_(pos) {}
    iterator() {}
  iterator(const iterator &o) :n_(o.n_), pos_(o.pos_) {}    
    const uint64 & Key() { return n_->Key(pos_); }
    V & Value() { return n_->Value(pos_); }
    iterator & operator ++() {
      // Make going forward from r_end() work
      if (pos_ == -2) {
	*this = root_to_begin(n_);
	return *this;
      }
      pos_++;
      if (pos_ < n_->count_) return *this;
      UpAndAdvance();
      AdvanceToTerminal();
      return *this;
    }
    iterator & operator --() {
      // Make going back from end() work
      if (pos_ == -1) {
	*this = root_to_r_begin(n_);
	return *this;
      }
      pos_--;
      if (pos_ >= 0) return *this;
      UpAndRetreat();
      RetreatToTerminal();
      return *this;
    }
    // Move to the node above the current node and point
    // to its next child.  May point past the end of the parent node.
    void UpAndAdvance() {
      int32 which = n_->which_child_;      
      // If at the root, point to end
      if (which == -1) {
	pos_ = -1;
	return; }
      n_ = n_->Parent();
      pos_ = which + 1;
    }

    // Move to the node above the current node and point
    // to its next child.  May point to position -1.
    void UpAndRetreat() {
      int32 which = n_->which_child_;      
      // If at the root, point to end
      if (which == -1) {
	pos_ = -2;
	return; }
      n_ = n_->Parent();
      pos_ = which - 1;
    }
    
    // Iterator starts out pointing to a position in a nonterminal.
    // Go to first key >= the start of this range, or to end()
    void AdvanceToTerminal() {
      while (1) {
	if (pos_ == -1) return;
	else if (pos_ == ARITY) UpAndAdvance();
	else if (n_->Child(pos_)) break;
	else pos_++;
      }
      n_ = n_->Child(pos_);
      while (!n_->IsTerminal()) {
	pos_ = 0;
	while (!n_->Child(pos_)) pos_++;
	CHECK(pos_ < ARITY);
	n_ = n_->Child(pos_);
      }
      pos_ = 0;
      CHECK(n_->count_ > 0);
    }

    // Iterator starts out pointing to a position in a nonterminal.
    // Go to last key < the end of this range, or to r_end()
    void RetreatToTerminal() {
      while (1) {
	if (pos_ == -2) return;
	else if (pos_ == -1) UpAndRetreat();
	else if (n_->Child(pos_)) break;
	else pos_--;
      }
      n_ = n_->Child(pos_);
      while (!n_->IsTerminal()) {
	pos_ = ARITY - 1;
	while (!n_->Child(pos_)) pos_--;
	CHECK(pos_ >= 0);
	n_ = n_->Child(pos_);
      }
      pos_ = n_->count_ -1;
      CHECK(n_->count_ > 0);
    }

    bool operator ==(const iterator & o) {
      return n_ == o.n_ && pos_ == o.pos_;
    }
    bool operator !=(const iterator &o) {
      return !(*this == o);
    }
    void erase(IntegerMap *m) {
      CHECK(n_->IsTerminal());
      CHECK(pos_ >= 0);
      CHECK(pos_ < n_->count_);
      for (int64 i=pos_; i+1 < n_->count_; i++) {
	n_->Key(i) = n_->Key(i+1);
	n_->Value(i) = n_->Value(i+1);
      }
      for (Node * p = n_; p; p = p->Parent()) p->count_--;
      CheckUnderfull(m);
    }

    // If the count is zero, delete it. 
    // LATER: If the count is too small, collapse into a terminal.
    void CheckUnderfull(IntegerMap *m) {
      while(1) {
	Node * parent = n_->Parent();
	if (!parent) return;
	if (n_->count_ == 0) {
	  parent->Child(n_->which_child_) = NULL;
	  m->Free(n_);
	  n_ = parent;
	} else {
	  return;
	}
      }
    }
  };  
  // end is root_, -1
  // r_end is root_, -2
  // begin and r_begin are dereferencable mostly unless they are end or r_end
  iterator begin() {
    return root_to_begin(root_);
  }
  static iterator root_to_begin(Node * root) {
    if (root->IsTerminal()) {
      if (root->count_ == 0) return root_to_end(root);
      return iterator(root, 0);
    }
    iterator ret(root, 0);
    ret.AdvanceToTerminal();
    return ret;
  }
  iterator end() {
    return root_to_end(root_);
  }
  static iterator root_to_end(Node * root) {
    return iterator(root, -1);
  }
  iterator r_begin() {
    return root_to_r_begin(root_);
  }
  static iterator root_to_r_begin(Node * root) {
    if (root->IsTerminal()) {
      if (root->count_ == 0) return root_to_r_end(root);
      return iterator(root, root->count_ - 1);
    }
    iterator ret(root, ARITY-1);
    ret.RetreatToTerminal();
    return ret;
  }
  iterator r_end() {
    return root_to_r_end(root_);
  }
  static iterator root_to_r_end(Node * root) {
    return iterator(root, -2);
  }

  iterator find(uint64 k) {
    Node * n = root_;
    while (1) {
      if (n->IsTerminal()) {
	// linearly scan for the position to insert this key.
	for (int64 i=0; i<n->count_; i++) {
	  if (n->Key(i) >= k) {
	    if (n->Key(i) == k) return iterator(n, i);
	    return end();
	  }
	}
	return end();
      }
      // It's a nonterminal node
      uint64 which = TopBits(k, n->lg_range_);
      n = n->Child(which);
      if (!n) return end();
    }
  }
  
  // return iterator pointing to the first key >= k
  iterator lower_bound(uint64 k) {
    Node * n = root_;
    while (1) {
      if (n->IsTerminal()) {
	// linearly scan for the position to insert this key.
	for (int64 i=0; i<n->count_; i++) 
	  if (n->Key(i) >= k) return iterator(n, i);
	iterator ret(n, n->count_-1);
	return ++ret;
      }
      // It's a nonterminal node
      uint64 which = TopBits(k, n->lg_range_);
      if (n->Child(which)) {
	n = n->Child(which);
	continue;
      }      
      iterator ret(n, which);
      ret.AdvanceToTerminal();
      return ret;
    }
  }  

  void erase(int64 key) {
    iterator look = find(key);
    if (look == end()) return;
    look.erase(this);
  }
  void erase(iterator i) {
    i.erase(this);
  }
  
  // a node represents a range of keys [a * ARITY^b, (a+1) * ARITY^b )
  //    for some a and b.  a and b are not represented in the node.
  // a node is either a terminal or an interior node.
  // if it's a terminal, data_ contains sorted key-value pairs. 
  // it it's an interior node, data_ contains pointers to child nodes. 
  struct Node {
    uint64 data_[ARITY];
    int64 count_;  // the number of existing keys under this node. 
    // top bits are a pointer to the parent. 
    // we steal the bottom bit as a flag that indicates terminal node.
    uint64 parent_; 
    int32 which_child_;
    int32 lg_range_; 
    
    inline bool IsTerminal() { return parent_ & 1ull; }
    inline Node * Parent() { return (Node *) (parent_ & (~1ull) ); }
    inline void ToggleTerminal() {parent_ ^= 1; }
    inline KV & KeyValue(int32 i) {return ((KV *)data_)[i];}
    inline uint64 & Key(int32 i) { 
      return *(data_ + i * (1 + sizeof(V) / sizeof(uint64)));
      // KeyValue(i).first;
    }
    inline V & Value(int32 i) { 
      return *((V *) (data_ + 1 + i * (1 + sizeof(V) / sizeof(uint64))));
      // KeyValue(i).second;
    }
    inline Node *& Child(int32 i) { return (Node *&)(data_[i]);}

    // is a terminal full
    inline bool IsFull() { 
      return ((count_+1) * sizeof(KV) > ARITY * sizeof(uint64)); }    

    // create an empty terminal node with a given parent. 
  Node(Node * parent, int32 which_child, int32 lg_range)
  : count_(0), 
      parent_((uint64)parent | 1ull), 
      which_child_(which_child),
      lg_range_(lg_range) {}
    ~Node() {
      if (!IsTerminal()) {
	for (uint64 i=0; i<ARITY; i++) delete Child(i);
      }
    }

    void Print(string prefix) {
      if (IsTerminal()) {
	for (uint64 i=0; i<count_; i++) 
	  cout << prefix << " " << Key(i) << " " << Value(i) << endl;
      } else {
	cout << prefix << " ." << endl;
	for (uint64 i=0; i<ARITY; i++) 
	  if (Child(i)) Child(i)->Print(prefix + char('0'+i));
      }
    }

    /*
    void PrintBackwards() {
      iterator i = end();
      while (i != r_end()) {
	--i;
	if (i != end()) {
	  cout << "key:" << i.Key() << endl;
	}
      }
    }
    */

    // finds or inserts. 
    // lg_range is the binary log of the range of integers represented under 
    //   this node. 
    V & Ref(uint64 k, int64 lg_range, IntegerMap *m) {
      // cout << "Ref " << k << " " << lg_range << endl;
      if (IsTerminal()) {
	// linearly scan for the position to insert this key.
	uint64 pos = count_; 
	for (int64 i=0; i<count_; i++) {
	  if (Key(i) >= k) {
	    if (Key(i) == k) return Value(i); // Key found.  Just return. 
	    pos = i;
	    break;
	  }
	}
	if (!IsFull()) {
	  // there is room in the terminal.  Shove things over and insert. 
	  for (uint64 i=count_; i>pos ; i--) {
	    Key(i) = Key(i-1);
	    Value(i) = Value(i-1);
	    //KeyValue(i) = KeyValue(i-1);
	  }
	  Key(pos) = k;
	  Value(pos) = V();
	  for (Node *n = this; n; n = n->Parent()) n->count_++;
	  return Value(pos);
	}
	// make this node a non-terminal
	Node * tmp[ARITY]; // this will be the new array of children. 
	for (uint64 i=0; i<ARITY; i++) tmp[i] = 0;
	for (int64 i=0; i<count_; i++) {
	  uint64 which = TopBits(Key(i), lg_range);
	  Node *& child = tmp[which];
	  if (!child) child = m->AllocateNode(this, which, lg_range - LG_ARITY);
	  // child->KeyValue(child->count_) = KeyValue(i);
	  child->Key(child->count_) = Key(i);
	  child->Value(child->count_) = Value(i);
	  child->count_++;
	}
	for (uint64 i=0; i<ARITY; i++) Child(i) = tmp[i];
	ToggleTerminal();
      }
      // now it's a non-terminal
      uint64 which = TopBits(k, lg_range);
      if (!Child(which)) 
	Child(which) = m->AllocateNode(this, which, lg_range - LG_ARITY);
      return Child(which)->Ref(k, lg_range - LG_ARITY, m);
    }
  };

  Node * AllocateNode(Node *parent, int64 which, int64 lg_range) {
    // cout << "allocatenode " << a_ << endl;
    if (a_) return new (*a_) Node (parent, which, lg_range);
    return new Node(parent, which, lg_range);
  }
  void Free(Node *n) { if (!a_) delete n; }

  void Print() {
    cout << "lg_range = " << lg_range_ << endl;
    root_->Print("");
  }
  
  IntegerMap() {root_ = NULL;}
  void Init(uint64 range, TempArena *a){
    // cout << "init " << a << endl;
    a_ = a;
    // Computes lg_range_ based on splitting by ARITY at each level
    // and what is requested
    lg_range_ = 0;
    uint64 r = 1;
    while (r < range) {
      lg_range_ += LG_ARITY;
      r <<= LG_ARITY;
    }
    root_ = AllocateNode(NULL, -1, lg_range_);
  }
  ~IntegerMap() { if (!a_) delete root_; }

  Node * root_;
  uint64 lg_range_;
  TempArena *a_;
  
  V & operator [] (uint64 k) {
    return root_->Ref(k, lg_range_, this);
  }
  
  uint64 size() {
    return root_->count_;
  }
};

#endif
